// Copyright 2025 The XLS Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <memory>
#include <string>

#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "xls/common/fuzzing/fuzztest.h"
#include "absl/log/log.h"
#include "absl/strings/str_format.h"
#include "xls/common/status/matchers.h"
#include "xls/fuzzer/ir_fuzzer/ir_fuzz_domain.h"
#include "xls/fuzzer/ir_fuzzer/ir_fuzz_test_library.h"
#include "xls/ir/ir_matcher.h"
#include "xls/ir/node.h"
#include "xls/ir/op.h"
#include "xls/ir/package.h"
#include "xls/ir/verifier.h"
#include "xls/passes/reassociation_pass.h"

namespace m = ::xls::op_matchers;
namespace xls {
namespace {

using testing::AnyOf;
using testing::Each;
// Perform tests on the Package object which contains the IR. This is a general
// test that just verifies if the Package object is valid.
void VerifyIrFuzzPackage(std::shared_ptr<Package> p) {
  XLS_ASSERT_OK(VerifyPackage(p.get()));
  VLOG(3) << "IR Fuzzer-2: IR:" << "\n" << p->DumpIr() << "\n";
}
// Use of gtest FUZZ_TEST to randomly generate IR while being compatible with
// Google infrastructure. The IrFuzzTest function is called and represents the
// main test logic. A domain is specified to define the range of possible values
// that the FuzzProgram protobuf can have when generating random values.
FUZZ_TEST(IrFuzzTest, VerifyIrFuzzPackage).WithDomains(IrFuzzDomain());

void BuilderRestrictsOp(FuzzPackage fuzz_package) {
  EXPECT_THAT(fuzz_package.p->functions().front()->nodes(),
              Each(AnyOf(
                  // always allowed
                  m::Param(), m::Literal(),
                  // Generated by type coercion
                  m::BitSlice(), m::ZeroExt(), m::SignExt(),
                  // For ret.
                  m::Tuple(),
                  // specifically allowed
                  m::Sub())))
      << fuzz_package.fuzz_program.DebugString();
}
FUZZ_TEST(IrFuzzTest, BuilderRestrictsOp)
    .WithDomains(FuzzPackageDomainBuilder()
                     .WithOnlyBitsOperations()
                     .WithOperations({Op::kSub})
                     .NoClz()
                     .NoCtz()
                     .Build());

// This test makes sure that the OptimizationPassChangesOutputs test works for
// this specific example.
TEST(IrFuzzTest, PassChangesOutputsWithBitsParam) {
  std::string proto_string = absl::StrFormat(
      R"(
        combine_list_method: LAST_ELEMENT_METHOD
        args_bytes: "\xf7\x80\x6a\x4a\x69\xe3\x12\x20\x0f\xbe\xab\x08\xcc\xeb\xf5\x14\xda\x6d\x08\xa5\x0c\x0e\xd0\xa5\x64\x02\xed\x35\x2c\xad\xa3\x23"
        fuzz_ops {
          param {
            type {
              bits {
                bit_width: 10
              }
            }
          }
        }
      )");
  ReassociationPass pass;
  XLS_ASSERT_OK(
      PassChangesOutputsWithProto(proto_string, /*arg_set_count=*/10, pass));
}

TEST(IrFuzzTest, PassChangesOutputsWithTwoBitsParams) {
  std::string proto_string = absl::StrFormat(
      R"(
        combine_list_method: LAST_ELEMENT_METHOD
        args_bytes: "\xa0\x02\xca\x71\xf7\x9b\x6e\xf7\x5d\xf6\xf8\xee\xe3\xa9\xc8\x2e\xe3\xf9\x52\x2d\xc5\x0f\x63\x05\x48\x7e\x25\x1d\x9d\xe3\x54\xa5\xeb\x53\xae\xc3\x6f\x1d\x8c\xc0\x47\xfd\x88\x57\x91\x3b\x0a\x3a\xeb\x3f\xa3\xf2\x0c\xa1\x31\x45\x80\x75\xf5\xc5\x51\xb7\xf9\x1c"
        fuzz_ops {
          param {
            type {
              bits {
                bit_width: 10
              }
            }
          }
        }
        fuzz_ops {
          param {
            type {
              bits {
                bit_width: 20
              }
            }
          }
        }
      )");
  ReassociationPass pass;
  XLS_ASSERT_OK(
      PassChangesOutputsWithProto(proto_string, /*arg_set_count=*/10, pass));
}

TEST(IrFuzzTest, PassChangesOutputsWithTupleParam) {
  std::string proto_string = absl::StrFormat(
      R"(
        combine_list_method: LAST_ELEMENT_METHOD
        args_bytes: "\x58\xa7\xe5\xf6\x60\x1b\xee\xac\x9a\x82\xe4\x4e\x95\xb3\xbc\x31\xa8\x61\x86\x97\x77\x90\xa3\x45\x6b\x50\x58\xf5\xde\xb8\xdc\xb2\x0d\x1c\xf4\x7f\x27\xde\x09\x60\x90\x85\xbb\x27\x4b\x0e\x18\xb1\xa9\xf1\x92\x17\x7d\x56\x90\x1a\x3d\x30\x4a\x78\x29\x2b\x21\x68"
        fuzz_ops {
          param {
            type {
              tuple {
                tuple_elements {
                  bits {
                    bit_width: 10
                  }
                }
                tuple_elements {
                  bits {
                    bit_width: 20
                  }
                }
              }
            }
          }
        }
      )");
  ReassociationPass pass;
  XLS_ASSERT_OK(
      PassChangesOutputsWithProto(proto_string, /*arg_set_count=*/10, pass));
}

TEST(IrFuzzTest, PassChangesOutputsWithArrayParam) {
  std::string proto_string = absl::StrFormat(
      R"(
        combine_list_method: LAST_ELEMENT_METHOD
        args_bytes: "\xe9\x32\x4f\xa7\x1d\xad\x40\x3b\xdc\x0f\xe7\x65\xb4\x2d\xac\x1d\x1d\x20\x35\x54\xc4\x49\x02\x28\x71\x3d\x43\x3c\x80\xba\xb5\xaa\x33\x9c\xa8\x8c\xe4\xa0\xd1\xba\x29\x65\xc2\x73\x57\xcd\x94\x18"
        fuzz_ops {
          param {
            type {
              array {
                array_size: 2
                array_element {
                  bits {
                    bit_width: 10
                  }
                }
              }
            }
          }
        }
      )");
  ReassociationPass pass;
  XLS_ASSERT_OK(
      PassChangesOutputsWithProto(proto_string, /*arg_set_count=*/10, pass));
}

TEST(IrFuzzTest, PassChangesOutputsWithNestedTupleParam) {
  std::string proto_string = absl::StrFormat(
      R"(
        combine_list_method: LAST_ELEMENT_METHOD
        args_bytes: "\xbd\x0f\x24\xdf\xbc\x9d\x56\x08\xf1\xf0\x8d\xc4\x65\x04\x08\x3e\xc0\x2e\x74\xdc\x23\xaf\x3f\x6d\xca\x59\x9a\x22\x8d\x42\x08\x7e\x42\x5b\x71\xf2\x42\x92\x2f\xfa\xaf\x73\x7d\x33\x3b\xed\x6a\xf9\x47\x7d\x58\x11\xe7\x4f\xec\xf4\x5d\xae\x53\xb3\x66\x5d\x69\x27"
        fuzz_ops {
          param {
            type {
              tuple {
                tuple_elements {
                  tuple {
                    tuple_elements {
                      bits {
                        bit_width: 10
                      }
                    }
                  }
                }
                tuple_elements {
                  bits {
                    bit_width: 20
                  }
                }
              }
            }
          }
        }
      )");
  ReassociationPass pass;
  XLS_ASSERT_OK(
      PassChangesOutputsWithProto(proto_string, /*arg_set_count=*/10, pass));
}

TEST(IrFuzzTest, PassChangesOutputsWithEmptyArrayParam) {
  std::string proto_string = absl::StrFormat(
      R"(
        combine_list_method: TUPLE_LIST_METHOD
        args_bytes: "\x4d\xa7\xa6\x21\x56\x9a\x0d\x42\xed\x0f\xc2\x38\xd3\xaa\x49\x67"
        fuzz_ops {
          param {
            type {
              array {
                array_size: 0
                array_element {
                  bits {
                    bit_width: 10
                  }
                }
              }
            }
          }
        }
      )");
  ReassociationPass pass;
  XLS_ASSERT_OK(
      PassChangesOutputsWithProto(proto_string, /*arg_set_count=*/10, pass));
}

TEST(IrFuzzTest, PassChangesOutputsWithParamAndLiteralTuple) {
  std::string proto_string = absl::StrFormat(
      R"(
        combine_list_method: TUPLE_LIST_METHOD
        args_bytes: "\x9c\xbe\x22\xda\x0f\x89\x3b\x2b\x01\x5a\xcc\x81\x12\x23\xea\x45\x7c\x14\xc7\x82\x47\xe5\x55\x54\xb5\xce\x45\x17\xed\xce\x20\xe2"
        fuzz_ops {
          param {
            type {
              tuple {
                tuple_elements {
                  bits {
                    bit_width: 10
                  }
                }
              }
            }
          }
        }
        fuzz_ops {
          literal {
            type {
              tuple {
                tuple_elements {
                  bits {
                    bit_width: 10
                  }
                }
              }
            }
          }
        }
      )");
  ReassociationPass pass;
  XLS_ASSERT_OK(
      PassChangesOutputsWithProto(proto_string, /*arg_set_count=*/10, pass));
}

}  // namespace
}  // namespace xls
